import sys
import os
import shutil
import cv2
from PyQt6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QHBoxLayout, QListWidget, QLabel, QPushButton, QSpacerItem, QSizePolicy
)
from PyQt6.QtGui import QPixmap, QImage
from PyQt6.QtCore import Qt

VIDEO_DIR = "output"
OUTPUT_FOLDER_PREFIX = "outputd_"
IMAGE_COUNT = 4  # P0~P3

class VideoFrameViewer(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("视频帧浏览器")
        self.resize(1800, 900)

        main_layout = QHBoxLayout(self)

        # L区：视频列表
        self.video_list = QListWidget()
        self.video_list.setFixedWidth(480)  # 30%
        self.video_list.itemClicked.connect(self.load_video)
        self.video_list.currentRowChanged.connect(self.load_video_by_row)
        main_layout.addWidget(self.video_list)

        # M区：显示四张图片序号
        self.seq_list = QListWidget()
        self.seq_list.setFixedWidth(320)  # 20%
        main_layout.addWidget(self.seq_list)

        # R区：四个区域纵向分布
        self.r_layout = QVBoxLayout()
        self.r_widgets = []

        for i in range(IMAGE_COUNT):
            r_item_layout = self.create_r_item(i)
            self.r_layout.addLayout(r_item_layout)
            self.r_layout.addSpacerItem(QSpacerItem(20, 40, QSizePolicy.Policy.Minimum, QSizePolicy.Policy.Expanding))

        main_layout.addLayout(self.r_layout, 1)  # 50%
        self.setLayout(main_layout)

        # 当前视频信息
        self.current_video_path = None
        self.cap = None
        self.current_folder = None
        self.current_indices = [0] * IMAGE_COUNT
        self.old_indices = [0] * IMAGE_COUNT
        self.folder_image_files = []

        self.load_video_list()

    def load_video_by_row(self, row):
        if row < 0:
            return
        item = self.video_list.item(row)
        if item:
            self.load_video(item)
    def load_video_list(self):
        self.video_list.clear()
        for file in sorted(os.listdir(VIDEO_DIR)):
            if file.endswith(".mp4"):
                self.video_list.addItem(file)

    def create_r_item(self, idx):
        layout = QHBoxLayout()

        # 左按钮
        btn_left = QPushButton(f"L{idx}")
        btn_left.clicked.connect(lambda _, i=idx: self.change_frame(i, -1))
        layout.addWidget(btn_left)

        # 老文本区
        lbl_old = QLabel("000")
        lbl_old.setFixedWidth(50)
        lbl_old.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(lbl_old)

        # 图片显示
        lbl_image = QLabel()
        lbl_image.setFixedSize(300, 180)
        lbl_image.setAlignment(Qt.AlignmentFlag.AlignCenter)
        lbl_image.setStyleSheet("border: 1px solid black;")
        layout.addWidget(lbl_image)

        # 新文本区
        lbl_new = QLabel("000")
        lbl_new.setFixedWidth(50)
        lbl_new.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(lbl_new)

        # 右按钮
        btn_right = QPushButton(f"R{idx}")
        btn_right.clicked.connect(lambda _, i=idx: self.change_frame(i, 1))
        layout.addWidget(btn_right)

        # 保存按钮
        btn_save = QPushButton(f"S{idx}")
        btn_save.clicked.connect(lambda _, i=idx: self.save_frame(i))
        layout.addWidget(btn_save)

        self.r_widgets.append({
            "image": lbl_image,
            "old_text": lbl_old,
            "new_text": lbl_new
        })

        return layout

    def load_video(self, item):
        video_name = item.text()
        self.current_video_path = os.path.join(VIDEO_DIR, video_name)
        self.cap = cv2.VideoCapture(self.current_video_path)

        # 获取 outputd 文件夹
        base_name = os.path.splitext(video_name)[0]
        self.current_folder = os.path.join(OUTPUT_FOLDER_PREFIX, f"{base_name}_folder")

        # 获取文件夹图片文件，按序号排序
        self.folder_image_files = sorted(
            [f for f in os.listdir(self.current_folder) if f.endswith(".png")]
        )

        # 更新中间 M 区显示序号
        self.seq_list.clear()
        seq_numbers = []
        for f in self.folder_image_files:
            seq = os.path.splitext(f)[0]
            self.seq_list.addItem(seq)
            seq_numbers.append(int(seq))

        # 初始化 R 区显示
        for i in range(IMAGE_COUNT):
            if i < len(seq_numbers):
                idx = seq_numbers[i]
            else:
                idx = 0
            self.current_indices[i] = idx
            self.old_indices[i] = idx
            self.r_widgets[i]["old_text"].setText(f"{idx:03d}")
            self.r_widgets[i]["new_text"].setText(f"{idx:03d}")
            self.update_r_image(i)

    def change_frame(self, idx, delta):
        new_idx = self.current_indices[idx] + delta
        if new_idx < 0:
            new_idx = 0
        if self.cap:
            total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
            if new_idx >= total_frames:
                new_idx = total_frames - 1
        self.current_indices[idx] = new_idx
        self.r_widgets[idx]["new_text"].setText(f"{new_idx:03d}")
        self.update_r_image(idx)

    def update_r_image(self, idx):
        frame_idx = self.current_indices[idx]
        # 尝试读取 outputd 文件夹图片
        image_path = os.path.join(self.current_folder, f"{frame_idx:03d}.png")
        if os.path.exists(image_path):
            pixmap = QPixmap(image_path)
        else:
            # 读取视频帧
            if self.cap:
                self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
                ret, frame = self.cap.read()
                if ret:
                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    h, w, ch = frame.shape
                    bytes_per_line = ch * w
                    qimg = QImage(frame.data, w, h, bytes_per_line, QImage.Format.Format_RGB888)
                    pixmap = QPixmap.fromImage(qimg)
                else:
                    pixmap = QPixmap()
            else:
                pixmap = QPixmap()
        pixmap = pixmap.scaled(300, 180, Qt.AspectRatioMode.KeepAspectRatio)
        self.r_widgets[idx]["image"].setPixmap(pixmap)

    def save_frame(self, idx):
        old_idx = self.old_indices[idx]
        new_idx = self.current_indices[idx]

        old_path = os.path.join(self.current_folder, f"{old_idx:03d}.png")
        new_path = os.path.join(self.current_folder, f"{new_idx:03d}.png")

        # 删除老文件
        if os.path.exists(old_path):
            os.remove(old_path)

        # 保存当前显示图片
        pixmap = self.r_widgets[idx]["image"].pixmap()
        if pixmap:
            pixmap.save(new_path)

        # 更新老文本区
        self.old_indices[idx] = new_idx
        self.r_widgets[idx]["old_text"].setText(f"{new_idx:03d}")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    viewer = VideoFrameViewer()
    viewer.show()
    sys.exit(app.exec())
